# -*- coding: utf-8 -*-

from os import path

import torch
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import Dataset, DataLoader
import numpy as np
import matplotlib.pyplot as plt
import h5py

# computed from train set
RGB_MEAN = [0.5429, 0.5263, 0.4994]
RGB_STD = [0.2422, 0.2392, 0.2406]

def build_data(data_set='train', batch_size=4, transform=True):
    ''' build selected dataset
    '''
    if transform:

        if data_set == 'train':
            data_transform = transforms.Compose([
                            transforms.RandomResizedCrop(256, (.8, 1)),
                            transforms.RandomHorizontalFlip(),
                            transforms.RandomRotation(degrees=15),
                            transforms.ColorJitter(),
                            transforms.CenterCrop(224),
                            transforms.ToTensor(),
                            transforms.Normalize(RGB_MEAN, RGB_STD),
                            ])
            if_shuffle = True

        else:
            data_transform = transforms.Compose([
                            transforms.Resize((224, 224)),
                            transforms.ToTensor(),
                            transforms.Normalize(RGB_MEAN, RGB_STD)
                            ])
            if_shuffle = False

    else:
        data_transform = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor()
        ])  # only resize, for statistical comuptation on dataset
        if_shuffle = False

    data_dir = path.join('data', 'caltech-101', data_set)
    dataset = torchvision.datasets.ImageFolder(data_dir, transform=data_transform)
    dataloadder = torch.utils.data.DataLoader(
                    dataset, batch_size=batch_size, shuffle=if_shuffle)
    return dataset, dataloadder

def compute_rgb_mean_std(data_set='train', batch_size=4):
    ''' computer rgb mean and std from dataset
        
        reference:
          <https://discuss.pytorch.org/t/about-normalization-using-pre-trained-vgg16-networks/23560/6>
    '''
    _, loader = build_data(data_set, batch_size, transform=False)
    mean, std, sample_num = 0., 0., 0.
    for data, _ in loader:
        batch_num = data.size(0)
        data = data.view(batch_num, data.size(1), -1)
        mean += data.mean(2).sum(0)
        std += data.std(2).sum(0)
        sample_num += batch_num
    
    mean /= sample_num
    std /= sample_num
    return mean, std

def imshow(img, show=True):
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    if show:
        plt.show()

def code_to_class(code):
    ''' convert code to class name
    '''
    with open('data/categories.txt', 'r') as f:
        lines = f.readlines()
        return lines[code].rstrip('\n')

if __name__ == '__main__':

    # these codes are used to compute rgb mean & std
    
    valset, valloader = build_data(data_set='train')  # replace para if you want to compute on other part of dataset

    with open('data/categories.txt', 'w') as f:
        categories = valset.classes
        for category in categories:
            f.write(category + '\n')

    dataiter = iter(valloader)
    images, labels = dataiter.next()
    print(images.size())
    imshow(torchvision.utils.make_grid(images[:, :]))
    print('GroundTruth:', ' '.join('%s\t'% code_to_class(labels[j].numpy()) for j in range(4)))

    mean, std = compute_rgb_mean_std(batch_size=16)
    print(mean)
    print(std)
